import json
import os
from tqdm import tqdm
from src.dynamic_quantizer import DynamicQuantizer
from src.structure_quantizer import StructureQuantizer

archive_path = '/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Generation/archive/success'


if __name__ == '__main__':
    structure_processor = StructureQuantizer(structure_vocab_size=2048, model_path="/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Quantizer/checkpoints/AE.pt", cluster_dir="/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Quantizer/checkpoints", batch_size=1)
    dynamic_processor = DynamicQuantizer(encoder_path="/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Quantizer/checkpoints/dynamic_encoder.ckpt", kmeans_path="/home/zhangqiang/userdata/wangzeyuan/Dynamic/Dynamic.Quantizer/checkpoints/kmeans.pkl")

    with open('./output/test.jsonl', 'a') as f_out:
        protein_list = os.listdir(archive_path)
        for protein in tqdm(protein_list, total=len(protein_list)):
            structure_sequence = structure_processor(os.path.join(archive_path, protein, 'structure.pdb'), return_residue_seq=True)['structure.pdb']
            structure, sequence = structure_sequence['struct'], structure_sequence['seq']
            dynamic = dynamic_processor(os.path.join(archive_path, protein, 'dynamic.pdb'))

            f_out.write(json.dumps({'name': protein, 'sequence': sequence, 'structure': structure, 'dynamic': dynamic}) + '\n')


